# -*- coding: utf-8 -*-
from time import sleep
import rospy
from geometry_msgs.msg import Twist,PoseStamped
import sys
import tty
import termios
from std_msgs.msg import String
import time
from nav_msgs.msg import Odometry

# control sequence
control = ['OFFBOARD', 'ARM', 'HOVER']

# Per-UAV state will be stored in lists indexed by 0..N-1 corresponding to uav ids a..b
multi_cmd_vel_flu_pub = []
multi_cmd_pub = []
multi_cmd_height_sub = []
icounts = []
hover_controls = []

def height_callback(data, idx):
    """Handle height (odometry) message for UAV at index idx (corresponding to uav_id = start_id + idx).
    Uses per-UAV counters and publishers."""
    global icounts, hover_controls
    icounts[idx] += 1
    twist = Twist()
    current_height = data.pose.pose.position.z
    print(f"UAV idx={idx} height={current_height}")

    if current_height >= 0.8 and (not hover_controls[idx]):
        print(f'UAV idx={idx} reached hover height, stopping vertical motion')
        hover_controls[idx] = True
        cmd = control[2]
        twist.linear.z = 0.0
        # publish stop vertical velocity and HOVER command
        try:
            multi_cmd_vel_flu_pub[idx].publish(twist)
            multi_cmd_pub[idx].publish(cmd)
        except Exception as e:
            print('Publish error:', e)
        # If all UAVs have hovered, shutdown
        if all(hover_controls):
            rospy.signal_shutdown("All UAVs reached target height, stopping program.")
            print("Program exiting after all UAVs hover.")
    else:
        # initial command sequence per-uav (mimic previous behavior)
        cnt = icounts[idx]
        if cnt == 1:
            twist.linear.z = 0.8
            try:
                multi_cmd_vel_flu_pub[idx].publish(twist)
            except Exception as e:
                print('Publish error:', e)
        elif cnt == 2:
            twist.linear.z = 0.8
            cmd = control[0]
            try:
                multi_cmd_vel_flu_pub[idx].publish(twist)
                multi_cmd_pub[idx].publish(cmd)
            except Exception as e:
                print('Publish error:', e)
            sleep(1)
        elif cnt == 3:
            twist.linear.z = 0.8
            cmd = control[1]
            try:
                multi_cmd_vel_flu_pub[idx].publish(twist)
                multi_cmd_pub[idx].publish(cmd)
            except Exception as e:
                print('Publish error:', e)

if __name__ == "__main__":

    settings = termios.tcgetattr(sys.stdin)

    # Usage: python3 hover.py <start_id> <end_id> [multirotor_type] [control_type]
    # Example: python3 hover.py 1 3 iris vel
    if len(sys.argv) < 3:
        print("Usage: python3 hover.py <start_id> <end_id> [multirotor_type] [control_type]")
        sys.exit(1)

    try:
        start_id = int(sys.argv[1])
        end_id = int(sys.argv[2])
    except Exception as e:
        print('Invalid start/end ids:', e)
        sys.exit(1)

    if start_id > end_id:
        start_id, end_id = end_id, start_id

    multirotor_type = sys.argv[3] if len(sys.argv) > 3 else 'iris'
    control_type = sys.argv[4] if len(sys.argv) > 4 else 'vel'

    cmd = String()
    twist = Twist()

    node_name = f"hover_{start_id}_{end_id}"
    rospy.init_node(node_name)
    rate = rospy.Rate(30)

    # prepare arrays for N UAVs
    uav_ids = list(range(start_id, end_id + 1))
    N = len(uav_ids)
    multi_cmd_vel_flu_pub = [None] * N
    multi_cmd_pub = [None] * N
    multi_cmd_height_sub = [None] * N
    icounts = [0] * N
    hover_controls = [False] * N

    if control_type == 'vel':
        for idx, uav_id in enumerate(uav_ids):
            # publishers for each UAV
            multi_cmd_vel_flu_pub[idx] = rospy.Publisher('/xtdrone/' + multirotor_type + '_' + str(uav_id) + '/cmd_vel_flu', Twist, queue_size=10)
            multi_cmd_pub[idx] = rospy.Publisher('/xtdrone/' + multirotor_type + '_' + str(uav_id) + '/cmd', String, queue_size=10)

            # subscribe to Odometry for height monitoring
            sub_topic = f"/{multirotor_type}_{uav_id}/Odometry"
            try:
                rospy.Subscriber(sub_topic, Odometry, height_callback, callback_args=idx, queue_size=10)
            except Exception as e:
                print(f"Subscriber create failed for {sub_topic}: {e}")

    print(f"Started hover controller for UAVs {start_id}..{end_id} (type={multirotor_type}, control={control_type})")

    while not rospy.is_shutdown():
        rate.sleep()